package com.cwj.algorithm.uf;

import java.util.Arrays;

/**
 * @author chenwujie
 * @date 2020-12-22 09:38
 */
public class UnionFindDisjointSet {
    /**
     * 下标代表元素，值代表所在组。
     */
    private final int[] eleAndGroup;
    /**
     * 代表代表根元素，值代表该树深度
     */
    private final int[] rank;

    public UnionFindDisjointSet(int size) {
        this.eleAndGroup = new int[size];
        this.rank = new int[size];

        for (int i = 0; i < size; i++) {
            eleAndGroup[i] = i;
        }
        Arrays.fill(rank, 1);
    }

    /**
     * 查询两元素是否在同一分组中
     *
     * @param x1
     * @param x2
     * @return
     */
    public boolean connected(int x1, int x2) {
        int r1 = find(x1);
        int r2 = find(x2);
        return r1 == r2;
    }

    /**
     * 查询元素 x 所在分组
     *
     * @param x
     * @return
     */
    public int find(int x) {
        return x == eleAndGroup[x] ? x : (eleAndGroup[x] = find(eleAndGroup[x]));
    }

    /**
     * 合并元素x1、x2所在分组的元素
     *
     * @param x1
     * @param x2
     */
    public void union(int x1, int x2) {
        int r1 = find(x1);
        int r2 = find(x2);
        if (r1 != r2) {
            // 合并
            if (rank[r1] <= rank[r2]) {
                eleAndGroup[r1] = r2;
            } else {
                eleAndGroup[r2] = r1;
            }

            // 当两树深度不相等时，合并后的深度不变，
            // 当两树深度相等时，合并后的树深度+1
            if(rank[r1] == rank[r2]){
                rank[r2]++;
            }
        }
    }
}
